import pickle
import numpy as np
import gradio as gr
from PIL import Image

def load_optimal_knn_model(model_path="best_knn_model.pkl"):
    """加载保存的最优KNN模型"""
    try:
        with open(model_path, "rb") as file:
            optimal_model = pickle.load(file)
        return optimal_model
    except FileNotFoundError:
        print(f"错误：未找到模型文件 {model_path}")
        print("请先运行train_optimal_knn.py生成模型文件")
        exit(1)

# 加载最优KNN模型
optimal_knn = load_optimal_knn_model()

def preprocess_and_predict(image):
    """预处理输入图像并使用最优KNN模型预测数字"""
    if image is None:
        return "请先绘制一个数字"
    
    # 确保图像是PIL格式
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)
    
    # 图像预处理：转为灰度图、调整尺寸、归一化
    gray_image = image.convert("L")
    resized_image = gray_image.resize((8, 8), Image.Resampling.LANCZOS)
    image_array = np.array(resized_image)
    
    # 反转颜色（匹配数据集的黑白对比度）并归一化到0-16范围
    inverted_array = 255 - image_array
    normalized_array = (inverted_array / 255) * 16
    
    # 调整数组形状为模型输入格式
    flattened_array = normalized_array.flatten().reshape(1, -1)
    
    # 模型预测
    prediction = optimal_knn.predict(flattened_array)
    return f"预测结果：{prediction[0]}"

# 使用Sketchpad组件并设置画笔粗细（brush_radius=3表示较细的画笔）
interface = gr.Interface(
    fn=preprocess_and_predict,
    inputs=gr.Sketchpad(
        height=400, 
        width=400, 
        label="请在此绘制数字",
        brush_radius=2 # 画笔粗细，数值越小越细（建议3-5之间）
    ),
    outputs="text",
    title="KNN手写数字识别",
    description="请在画板上绘制0-9的数字，系统将自动识别"
)

if __name__ == "__main__":
    interface.launch(share=False)